{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 注解区域"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import relay\n",
    "from tvm.relay.op.annotation import compiler_begin, compiler_end\n",
    "\n",
    "def check_region(region_set, target, args, nodes, rets):\n",
    "    region = region_set.get_region(args[0])\n",
    "    assert region\n",
    "    assert target == region.target\n",
    "    assert set(args) == set(region.args)\n",
    "    assert set(nodes) == set(region.nodes)\n",
    "    assert set(rets) == set(region.rets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试区域集创建者菱形"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fn (%data: Tensor[(10, 10), float32]) {\n",
      "  %0 = annotation.compiler_begin(%data, compiler=\"test_target\");\n",
      "  %1 = abs(%0);\n",
      "  %2 = annotation.compiler_end(%1, compiler=\"test_target\");\n",
      "  %3 = annotation.compiler_begin(%2, compiler=\"test_target\");\n",
      "  %4 = nn.relu(%3);\n",
      "  %5 = annotation.compiler_end(%4, compiler=\"test_target\");\n",
      "  %6 = annotation.compiler_end(%1, compiler=\"test_target\");\n",
      "  %7 = annotation.compiler_begin(%6, compiler=\"default\");\n",
      "  %8 = tanh(%7);\n",
      "  %9 = annotation.compiler_end(%8, compiler=\"default\");\n",
      "  %10 = annotation.compiler_begin(%5, compiler=\"test_target\");\n",
      "  %11 = annotation.compiler_begin(%9, compiler=\"test_target\");\n",
      "  %12 = add(%10, %11);\n",
      "  annotation.compiler_end(%12, compiler=\"test_target\")\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "data = relay.var(\"data\", shape=(10, 10))\n",
    "cb_1 = compiler_begin(data, \"test_target\")\n",
    "O_1 = relay.abs(cb_1)\n",
    "ce_1 = compiler_end(O_1, \"test_target\")\n",
    "ce_2 = compiler_end(O_1, \"test_target\")\n",
    "cb_2 = compiler_begin(ce_1, \"test_target\")\n",
    "O_2 = relay.nn.relu(cb_2)\n",
    "ce_3 = compiler_end(O_2, \"test_target\")\n",
    "cb_d = compiler_begin(ce_2, \"default\")\n",
    "X = relay.tanh(cb_d)\n",
    "ce_d = compiler_end(X, \"default\")\n",
    "cb_3 = compiler_begin(ce_3, \"test_target\")\n",
    "cb_4 = compiler_begin(ce_d, \"test_target\")\n",
    "O_3 = relay.add(cb_3, cb_4)\n",
    "ce_4 = compiler_end(O_3, \"test_target\")\n",
    "diamond = relay.Function([data], ce_4)\n",
    "print(diamond)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "region_set = relay.analysis.AnnotatedRegionSet(\n",
    "    diamond, relay.op.get(\"annotation.compiler_begin\"), relay.op.get(\"annotation.compiler_end\")\n",
    ")\n",
    "assert len(region_set) == 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[0;31mInit signature:\u001b[0m \u001b[0mrelay\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0manalysis\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAnnotatedRegionSet\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mexpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mregion_begin_op\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mregion_end_op\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mDocstring:\u001b[0m      Class to represent a relay expression split into regions.\n",
      "\u001b[0;31mInit docstring:\u001b[0m\n",
      "Construct regions from an expression.\n",
      "\n",
      "Parameters\n",
      "----------\n",
      "expr : tvm.relay.Expr\n",
      "    The expression from which to construct the regions.\n",
      "region_begin_op : tvm.ir.Op\n",
      "    The region begin annotation.\n",
      "region_end_op : tvm.ir.Op\n",
      "    The region end annotation.\n",
      "\u001b[0;31mFile:\u001b[0m           /media/pc/data/lxw/ai/tvm/python/tvm/relay/analysis/annotated_regions.py\n",
      "\u001b[0;31mType:\u001b[0m           type\n",
      "\u001b[0;31mSubclasses:\u001b[0m     "
     ]
    }
   ],
   "source": [
    "relay.analysis.AnnotatedRegionSet?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "check_region(\n",
    "    region_set,\n",
    "    \"test_target\",\n",
    "    [cb_1],\n",
    "    [cb_1, O_1, ce_1, ce_2],\n",
    "    [ce_1, ce_2],\n",
    ")\n",
    "check_region(\n",
    "    region_set,\n",
    "    \"test_target\",\n",
    "    [cb_2],\n",
    "    [cb_2, O_2, ce_3],\n",
    "    [ce_3],\n",
    ")\n",
    "check_region(\n",
    "    region_set,\n",
    "    \"default\",\n",
    "    [cb_d],\n",
    "    [cb_d, X, ce_d],\n",
    "    [ce_d],\n",
    ")\n",
    "check_region(\n",
    "    region_set,\n",
    "    \"test_target\",\n",
    "    [cb_3, cb_4],\n",
    "    [cb_3, cb_4, O_3, ce_4],\n",
    "    [ce_4],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 测试区域集创建器合并"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fn (%data: Tensor[(10, 10), float32]) {\n",
      "  %0 = annotation.compiler_begin(%data, compiler=\"test_target\");\n",
      "  %1 = abs(%0);\n",
      "  %2 = nn.relu(%1);\n",
      "  %3 = annotation.compiler_end(%2, compiler=\"test_target\");\n",
      "  %4 = annotation.compiler_end(%1, compiler=\"test_target\");\n",
      "  %5 = annotation.compiler_begin(%4, compiler=\"default\");\n",
      "  %6 = tanh(%5);\n",
      "  %7 = annotation.compiler_end(%6, compiler=\"default\");\n",
      "  %8 = annotation.compiler_begin(%3, compiler=\"test_target\");\n",
      "  %9 = annotation.compiler_begin(%7, compiler=\"test_target\");\n",
      "  %10 = add(%8, %9);\n",
      "  %11 = add(%8, %9);\n",
      "  %12 = (%10, %11);\n",
      "  annotation.compiler_end(%12, compiler=\"test_target\")\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "data = relay.var(\"data\", shape=(10, 10))\n",
    "cb_1 = compiler_begin(data, \"test_target\")\n",
    "O_1 = relay.abs(cb_1)\n",
    "ce_2 = compiler_end(O_1, \"test_target\")\n",
    "O_2 = relay.nn.relu(O_1)\n",
    "ce_3 = compiler_end(O_2, \"test_target\")\n",
    "cb_d = compiler_begin(ce_2, \"default\")\n",
    "X = relay.tanh(cb_d)\n",
    "ce_d = compiler_end(X, \"default\")\n",
    "cb_3 = compiler_begin(ce_3, \"test_target\")\n",
    "cb_4 = compiler_begin(ce_d, \"test_target\")\n",
    "O_3 = relay.add(cb_3, cb_4)\n",
    "O_4 = relay.add(cb_3, cb_4)\n",
    "O_5 = relay.Tuple([O_3, O_4])\n",
    "ce_4 = compiler_end(O_5, \"test_target\")\n",
    "merged = relay.Function([data], ce_4)\n",
    "print(merged)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "region_set = relay.analysis.AnnotatedRegionSet(\n",
    "    merged, relay.op.get(\"annotation.compiler_begin\"), relay.op.get(\"annotation.compiler_end\")\n",
    ")\n",
    "assert len(region_set) == 3\n",
    "check_region(\n",
    "    region_set,\n",
    "    \"test_target\",\n",
    "    [cb_1],\n",
    "    [cb_1, O_1, O_2, ce_2, ce_3],\n",
    "    [ce_2, ce_3],\n",
    ")\n",
    "check_region(\n",
    "    region_set,\n",
    "    \"default\",\n",
    "    [cb_d],\n",
    "    [cb_d, X, ce_d],\n",
    "    [ce_d],\n",
    ")\n",
    "check_region(\n",
    "    region_set,\n",
    "    \"test_target\",\n",
    "    [cb_3, cb_4],\n",
    "    [cb_3, cb_4, O_3, O_4, O_5, ce_4],\n",
    "    [ce_4],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
